package com.qomosoft.mediacodectest;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;

import java.io.IOException;
import java.nio.ByteBuffer;

/**
 * Created by robin on 17/11/2017.
 */

public class QMAVCEncoder extends QMMediaCodec {
    private QMAVCEncoderCallback mCallback;
    private static long GENERATE_INDEX = 0;
    public byte[] configbyte;
    private int m_width;
    private int m_height;

    public interface QMAVCEncoderCallback {
        void onDataEncode(ByteBuffer buffer, MediaCodec.BufferInfo bufferInfo);
    }

    public QMAVCEncoder(int width, int height, int frameRate, QMAVCEncoderCallback callback) {
        this.mCallback = callback;
        m_width = width;
        m_height = height;
        if (null == mMediaCodec) {
            try {
                mMediaCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC);
                MediaFormat mediaFormat = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, MainActivity.PREVIEW_WIDTH, MainActivity.PREVIEW_HEIGHT);
                mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, 5 * width * height);
                mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar);
                mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1);
                mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, frameRate);
                mMediaCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
                mMediaCodec.start();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public void encode(byte[] data) {
        byte[] input = null;

        if (null != data) {
            byte[] yuv420sp = new byte[m_width * m_height * 3 / 2];
            //把待编码的视频帧转换为YUV420格式
            NV21ToNV12(data, yuv420sp, m_width, m_height);
            input = yuv420sp;
        }
        if (input != null) {
            try {
                //编码器输入缓冲区
                ByteBuffer[] inputBuffers = mMediaCodec.getInputBuffers();
                //编码器输出缓冲区
                ByteBuffer[] outputBuffers = mMediaCodec.getOutputBuffers();
                int inputBufferIndex = mMediaCodec.dequeueInputBuffer(-1);
                if (inputBufferIndex >= 0) {
                    ByteBuffer inputBuffer = inputBuffers[inputBufferIndex];
                    inputBuffer.clear();
                    //把转换后的YUV420格式的视频帧放到编码器输入缓冲区中
                    inputBuffer.put(input);
                    mMediaCodec.queueInputBuffer(inputBufferIndex, 0, input.length, computePresentationTime(GENERATE_INDEX++), 0);
                }

                MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
                int outputBufferIndex = mMediaCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC);
                while (outputBufferIndex >= 0) {
                    //Log.i("AvcEncoder", "Get H264 Buffer Success! flag = "+bufferInfo.flags+",pts = "+bufferInfo.presentationTimeUs+"");
                    ByteBuffer outputBuffer = outputBuffers[outputBufferIndex];
//                    byte[] outData = new byte[bufferInfo.size];
//                    outputBuffer.get(outData);
//                    if (bufferInfo.flags == MediaCodec.BUFFER_FLAG_CODEC_CONFIG) {
//                        configbyte = new byte[bufferInfo.size];
//                        configbyte = outData;
//                    } else if (bufferInfo.flags == MediaCodec.BUFFER_FLAG_KEY_FRAME) {
//                        byte[] keyframe = new byte[bufferInfo.size + configbyte.length];
//                        System.arraycopy(configbyte, 0, keyframe, 0, configbyte.length);
//                        //把编码后的视频帧从编码器输出缓冲区中拷贝出来
//                        System.arraycopy(outData, 0, keyframe, configbyte.length, outData.length);
//
////                        outputStream.write(keyframe, 0, keyframe.length);
//                        if (mCallback != null) {
//                            mCallback.onDataEncode(keyframe, bufferInfo);
//                        }
//                    } else {
//                        //写到文件中
////                        outputStream.write(outData, 0, outData.length);
//                        if (mCallback != null) {
//                            mCallback.onDataEncode(outData, bufferInfo);
//                        }
//                    }
                    mCallback.onDataEncode(outputBuffer, bufferInfo);

                    mMediaCodec.releaseOutputBuffer(outputBufferIndex, false);
                    outputBufferIndex = mMediaCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC);
                }

            } catch (Throwable t) {
                t.printStackTrace();
            }
        }
    }

    private void NV21ToNV12(byte[] nv21,byte[] nv12,int width,int height){
        if(nv21 == null || nv12 == null)return;
        int framesize = width*height;
        int i = 0,j = 0;
        System.arraycopy(nv21, 0, nv12, 0, framesize);
        for(i = 0; i < framesize; i++){
            nv12[i] = nv21[i];
        }
        for (j = 0; j < framesize/2; j+=2)
        {
            nv12[framesize + j-1] = nv21[j+framesize];
        }
        for (j = 0; j < framesize/2; j+=2)
        {
            nv12[framesize + j] = nv21[j+framesize-1];
        }
    }

    /**
     * Generates the presentation time for frame N, in microseconds.
     */
    private long computePresentationTime(long frameIndex) {
        return 132 + frameIndex * 1000000 / MainActivity.PREVIEW_FRAMERATE;
    }
}
